#include "CogEvolver.h"

#undef max
#undef min
#include <limits>

// ------------------------------------------------------------------------------------------

CogProductionStats::CogProductionStats()
	: instanceCount(0),
	  instanceUsageMean(0),
	  instanceUsageStd(0),
	  instanceUsageMax(0),
	  instanceAgeMean(0),
	  instanceAgeStd(0),
	  instanceAgeMax(0),
	  ntSizeMean(0),
	  ntSizeStd(0),
	  ntSizeMin(numeric_limits<double>::max()),
	  ntSizeMax(0)
{
}

// ------------------------------------------------------------------------------------------

CogEvalThread::CogEvalThread(CogEvolver* _pEvolver, size_t _stopIndex)
: pEvolver(_pEvolver),
  stopIndex(_stopIndex)
{
}

void CogEvalThread::operator()()
{
	pEvolver->evalMutex.lock();
	while (pEvolver->currentEvalIndex < stopIndex) 
	{
		size_t processIndex = pEvolver->currentEvalIndex++;
		pEvolver->evalMutex.unlock();
		pEvolver->evaluate(processIndex);
		pEvolver->evalMutex.lock();
	}
	pEvolver->evalMutex.unlock();
}

// ------------------------------------------------------------------------------------------

void CogEvolver::initialize()
{
	currentGeneration = 0;
	pSolutions = &offspring;
	allocSolutions();
	pSolutions = &parents;
	allocSolutions();
	numSolutions = 0; // valid solutions
	cogSeedRandom(parameters.randomSeed); // seed random number generator
	grammar.initialize(this);
}

void CogEvolver::evolve()
{
	CogStopWatch stopwatch;
	// evolve only if maxEpochs > 0
	if (parameters.maxEpochs)
	{
		globalMessageLog.print(COGMT_INFO_SYSTEM, VERBOSITY_LOW, "Initializing... (with seed: %d)", parameters.randomSeed);
		// clear the stats
		setupStats();
		stopwatch.startTimer();
		// initialize population: produce maxPop solution candidates by instantiating from templates
		fillSolutions(parameters.maxPop);
		// record initialization time cost
		stopwatch.stopTimer();
		addClockStat(globalStats, stopwatch.getElapsedTime());
		// repeat for each generation
		while (currentGeneration < parameters.maxEpochs)
		{
			globalMessageLog.print(COGMT_INFO_SYSTEM, VERBOSITY_LOW, "Evolving generation %d out of %d (%6.2f%%)", currentGeneration+1, parameters.maxEpochs, static_cast<float>(currentGeneration+1)/parameters.maxEpochs*100);
			// reproduction and evaluation
			evolveStep();
			// early exit if condition met
			if (parameters.stopCondition & COGSC_PRIMARYZERO)
			{
				if ((*pSolutions)[0].hasPrimaryZero())
				{
					break;
				}
			}
		}
	}
	else // evaluate only
	{
		stopwatch.startTimer();
		evaluateSolutions();
		// display intermediate stats
		if (parameters.verbosity > VERBOSITY_LOW)
		{
			ostringstream strStream;
			strStream << "-- Solutions (evaluated): " << numSolutions << endl;
			strStream << setiosflags(ios::fixed) << setprecision(2);
			strStream << "BEST--";
			for (size_t s = 0; s < parameters.objectives.stages(); s++)
			{
				for (size_t t = 0; t < parameters.objectives.stageSize(s); t++)
				{
					strStream << " " << parameters.objectives.objectiveName(s, t).c_str() << ": " << (*pSolutions).front().result[s][t];
				}
			}
			globalMessageLog.printStream(COGMT_INFO_STATS, VERBOSITY_NORMAL, strStream);
		}
	}
	stopwatch.stopTimer();
	globalMessageLog.print(COGMT_INFO_STATS, VERBOSITY_NORMAL, "****** Total time: %f", stopwatch.getElapsedTime());
}

void CogEvolver::evaluateSolutions(size_t startIndex, size_t endIndex)
{
	boost::thread_group threads;
	currentEvalIndex = startIndex;
	for (size_t i = 0; i < parameters.evalThreads; i++)
	{
        threads.create_thread(CogEvalThread(this, endIndex));
	}
	threads.join_all();


	/*boost::thread_group threads;
	size_t blockSize = (endIndex - startIndex) / parameters.evalThreads;
	size_t blockRem = (endIndex - startIndex) % parameters.evalThreads;
	size_t stop = 0;
    for (size_t i = 0; i < parameters.evalThreads; i++)
	{
		size_t start = stop;
		stop = stop + blockSize + ((i < blockRem) ? 1 : 0);
        threads.create_thread(CogEvalJob(this, startIndex+start, startIndex+stop));
	}
    threads.join_all();*/
}

void CogEvolver::solutionsToStream(ostream &out) const
{
	out << "% Pool Id";
	for (size_t s = 0; s < parameters.objectives.stages(); s++)
	{
		for (size_t t = 0; t < parameters.objectives.stageSize(s); t++)
		{
			if (s > 0)
			{
				out << " " << parameters.objectives.objectiveName(s, t).c_str();
			}
			else // measure, non-objective
			{
				out << " (" << parameters.objectives.objectiveName(s, t).c_str() << ")";
			}
		}
	}
	out << " Solution" << endl;
	for (size_t i = 0; i < numSolutions; i++)
	{
		out << "% " << (*pSolutions)[i].pRule->lhs.pPool->getParent()->getLabel((*pSolutions)[i].pRule->lhs.pPool->getId())
			<< " " << (*pSolutions)[i].pRule->lhs.pPool->getLabel((*pSolutions)[i].pRule->lhs.id);
		for (size_t s = 0; s < parameters.objectives.stages(); s++)
		{
			for (size_t t = 0; t < parameters.objectives.stageSize(s); t++)
			{
				out << " " << (*pSolutions)[i].result[s][t];
			}
		}
		out << " ";
		(*pSolutions)[i].pRule->expandToStream(out, &grammar, true);
		out << endl;
	}
}

void CogEvolver::statsToStream(ostream &out) const
{
	out << "% epoch-time-tEval-s#-p#-useMn-useSd-useMx-ageMn-ageSd-ageMx-ntMn-ntSd-ntMx-ntMi";
	for (size_t i = 0; i < 2; i++)
	{
		out << "   ";
		for (size_t s = 0; s < parameters.objectives.stages(); s++)
		{
			for (size_t t = 0; t < parameters.objectives.stageSize(s); t++)
			{
				if (s > 0)
				{
					out << " " << parameters.objectives.objectiveName(s, t).c_str() << "[Best-Mn-Sd-Mi-Mx] ";
				}
				else // measure, non-objective
				{
					out << " (" << parameters.objectives.objectiveName(s, t).c_str() << "[Best-Mn-Sd-Mi-Mx] ";
				}
			}
		}
	}
	out << endl;
	for (size_t i = 0; i < parameters.maxEpochs; i++)
	{
		for (vector<double>::const_iterator g_iter = globalStats[i].begin(); g_iter != globalStats[i].end(); ++g_iter)
		{
			if ((*g_iter < 1.0e+100) && (*g_iter > -1.0e+100))
				out << *g_iter << " ";
			else
				out << 0 << " ";
		}
		out << "      ";
		for (vector<double>::const_iterator e_iter = eliteStats[i].begin(); e_iter != eliteStats[i].end(); ++e_iter)
		{
			if ((*e_iter < 1.0e+100) && (*e_iter > -1.0e+100))
				out << *e_iter << " ";
			else
				out << 0 << " ";
		}
		out << "      ";
		for (vector<double>::const_iterator o_iter = offspringStats[i].begin(); o_iter != offspringStats[i].end(); ++o_iter)
		{
			if ((*o_iter < 1.0e+100) && (*o_iter > -1.0e+100))
				out << *o_iter << " ";
			else
				out << 0 << " ";
		}
		out << "\n";
	}
}

// -----

void CogEvolver::evolveStep()
{
	CogStopWatch stopwatch, evalwatch;
	stopwatch.startTimer();
	// reproduce
	size_t numOffspring = 0;
	size_t numActualElite = (parameters.elitePop > numSolutions) ? numSolutions : parameters.elitePop;
	size_t fillPop = parameters.maxPop - numActualElite;
	size_t failedOffspring = 0;
	vector<CogSolution>* pChildren = (pSolutions == &parents) ? &offspring : &parents;
	if (!parameters.tourSize) // each parent produces one offspring (cycle if necessary)
	{
		size_t solIndex = 0;
		while (numOffspring < fillPop)
		{
			CogProduction* pNewRule = grammar.mutate((*pSolutions)[solIndex].pRule);
			if (pNewRule) // mutation successful
			{
				(*pChildren)[numOffspring+numActualElite].pRule = pNewRule;
				numOffspring++;
			}
			else
			{
				failedOffspring++;
			}
			if (++solIndex >= numSolutions)
			{
				solIndex = 0;
			}
		}
	}
	else
	if (parameters.tourSize >= parameters.maxPop) // best solution is sole reproducer
	{
		while (numOffspring < fillPop)
		{
			CogProduction* pNewRule = grammar.mutate((*pSolutions)[0].pRule);
			if (pNewRule) // mutation successful
			{
				(*pChildren)[numOffspring+numActualElite].pRule = pNewRule;
				numOffspring++;
			}
			else
			{
				failedOffspring++;
			}
		}
	}
	else // tournament based selection
	{
		while (numOffspring < fillPop)
		{
			size_t tourIndex = fillPop;
			for (size_t i = 0; i < parameters.tourSize; i++) // find the best solution in tournament
			{
				size_t index2 = cogRandomIndex(fillPop-i); // account for previously taken ones
				if (index2 < tourIndex)
				{
					tourIndex = index2;
				}
			}
			CogProduction* pNewRule = grammar.mutate((*pSolutions)[tourIndex].pRule);
			if (pNewRule) // mutation successful
			{
				(*pChildren)[numOffspring+numActualElite].pRule = pNewRule;
				numOffspring++;
			}
			else
			{
				failedOffspring++;
			}
		}
	}
	// copy elite
	copy(pSolutions->begin(), pSolutions->begin()+numActualElite, pChildren->begin());
	// delete parents
	for (size_t i = numActualElite; i < numSolutions; i++) // not applied to elite
	{
		(*pSolutions)[i].pRule->lhs.flag &= ~COGSF_STARTING;
		(*pSolutions)[i].pRule->unuse();
	}
	// clean up
	grammar.clearUnusedInstances();
	// flip parent/offspring storage
	pSolutions = pChildren;
	// elite+offspring are the new population
	numSolutions = numActualElite + numOffspring;
	evalwatch.startTimer();
	// evaluate the offspring
	evaluateSolutions(numActualElite, numSolutions);
	evalwatch.stopTimer();
	addClockStat(globalStats, evalwatch.getElapsedTime(), 2); 
	// record/display statistics for offspring
	computeAllSolutionStats(offspringStats, numActualElite, numSolutions, "NEW--- ");
	// sort the population
	rankSortSolutions();
	// make new productions valid choices for future offspring
	grammar.computeInstanceUsage();
	stopwatch.stopTimer();
	// record/display statistics for next elite
	computeAllSolutionStats(eliteStats, 0, numActualElite, "ELITE- ");
	// stats for the grammar (elite+offspring)
	computeGlobalStats(globalStats);
	// print stats of best solution
	printSolutionStats(0, "BEST--");
	addClockStat(globalStats, stopwatch.getElapsedTime()); // add eval time to init time
	addFailedStat(globalStats, failedOffspring);
	globalMessageLog.print(COGMT_INFO_STATS, VERBOSITY_HIGH, "------ Time: %f -- Failures: %d -- Random key: %u", stopwatch.getElapsedTime(), failedOffspring, cogRandomInteger());
	// increment generation counter
	currentGeneration++;
}

void CogEvolver::allocSolutions()
{
	pSolutions->resize(parameters.maxPop);
	for (size_t i = 0; i < pSolutions->size(); i++) // setup results storage
	{
		(*pSolutions)[i].result.resize(parameters.objectives.stages());
		for (size_t j = 0; j < parameters.objectives.stages(); j++)
		{
			(*pSolutions)[i].result[j].resize(parameters.objectives.stageSize(j));
		}
	}
}

void CogEvolver::fillSolutions(size_t wantedSize)
{
	if (grammar.startingTemplates.size() && (numSolutions < wantedSize))
	{
		size_t existing_sols = numSolutions;
		size_t tpl_idx = 0;
		size_t tpl_passes = 0;
		bool created_something = false;
		double probReuseOnCreate = parameters.probReuseOnCreate;
		if (parameters.buildMode & COGBM_BIGSTART) // temporarily switch to template-only instantiation
		{
			parameters.probReuseOnCreate = 0.0;
		}
		while (numSolutions < wantedSize)
		{
			CogProduction* pTemplate = grammar.startingTemplates[tpl_idx];
			if (!tpl_passes || !(pTemplate->lhs.flag & COGSF_CONST)) // const templates are only instantiated once as starting productions
			{
				instantiateSolution(pTemplate);
				created_something = true;
			}
			if (++tpl_idx == grammar.startingTemplates.size())
			{
				if (!created_something)
				{
					globalMessageLog.print(COGMT_INFO_SYSTEM, VERBOSITY_NONE, "Warning: Grammar is composed ONLY of terminal productions!");
					break;
				}
				tpl_idx = 0;
				created_something = false;
				tpl_passes++;
			}
		}
		if (parameters.buildMode & COGBM_BIGSTART) // un-switch
		{
			parameters.probReuseOnCreate = probReuseOnCreate;
		}
		CogStopWatch evalwatch;
		evalwatch.startTimer();
		evaluateSolutions(0, wantedSize); // evaluate all
		evalwatch.stopTimer();
		addClockStat(globalStats, evalwatch.getElapsedTime(), 2); 
		if (parameters.verbosity > VERBOSITY_NORMAL)
		{
			for (size_t i = 0; i < wantedSize; i++)
			{
				ostringstream strStream;
				if (i < existing_sols)
				{
					strStream << "==";
				}
				else
				{
					strStream << "++";
				}
				for (size_t s = 0; s < parameters.objectives.stages(); s++)
				{
					for (size_t t = 0; t < parameters.objectives.stageSize(s); t++)
					{
						strStream << " " << parameters.objectives.objectiveName(s, t).c_str() << ": " << (*pSolutions)[i].result[s][t];
					}
				}
				globalMessageLog.printStream(COGMT_INFO_STATS, VERBOSITY_HIGH, strStream);
			}
		}
	}
	// make new productions valid choices
	grammar.computeInstanceUsage();
}

/*
http://sci2s.ugr.es/docencia/doctobio/2002-6-2-DEB-NSGA-II.pdf
*/
void CogEvolver::rankSortSolutions()
{
	for (size_t stage = 1; stage < parameters.objectives.stages(); stage++) // start at 1 (0 is non-objective)
	{
		if (parameters.objectives.stageSize(stage) > 1)
		{
			list<size_t> F;
			vector<list<size_t>> S;
			vector<size_t> n;
			S.resize(numSolutions);
			n.resize(numSolutions);
			for (size_t p = 0; p < numSolutions; p++)
			{
				n[p] = 0;
				for (size_t q = 0; q < numSolutions; q++)
				{
					if (parameters.objectives.dominates((*pSolutions)[p], (*pSolutions)[q], stage))
					{
						S[p].push_back(q);
					}
					else
					if (parameters.objectives.dominates((*pSolutions)[q], (*pSolutions)[p], stage))
					{
						n[p] += 1;
					}
				}
				if (n[p] == 0)
				{
					(*pSolutions)[p].result[stage].rank = 0;
					F.push_back(p);
				}	
			}
			size_t R = 1;
			while (F.size())
			{
				list<size_t> Q;
				list<size_t>::iterator F_iter = F.begin();
				for (; F_iter != F.end(); F_iter++)
				{
					list<size_t>::iterator S_iter = S[*F_iter].begin();
					for (; S_iter != S[*F_iter].end(); S_iter++)
					{
						n[*S_iter] -= 1;
						if (n[*S_iter] == 0)
						{
							(*pSolutions)[*S_iter].result[stage].rank = R;
							Q.push_back(*S_iter);
						}
					}
				}
				R++;
				F = Q;
			}
		}
	}
	sort((*pSolutions).begin(), (*pSolutions).begin()+numSolutions, cogSolutionCompare);
}

// -----

void CogEvolver::setupStats()
{
	if (recordStats)
	{
		eliteStats.clear();
		offspringStats.clear();
		globalStats.clear();
		eliteStats.resize(parameters.maxEpochs);
		offspringStats.resize(parameters.maxEpochs);
		globalStats.resize(parameters.maxEpochs);
		for (size_t i = 0; i < parameters.maxEpochs; i++)
		{
			eliteStats[i].resize(5*parameters.objectives.objectiveSize(), 0.0f);
			offspringStats[i].resize(5*parameters.objectives.objectiveSize(), 0.0f);
			globalStats[i].resize(16, 0.0f);
		}
	}
}

void CogEvolver::computeAllSolutionStats(vector<vector<double>>& statMatrix, size_t startIndex, size_t stopIndex, const string& tag)
{
	if ((parameters.verbosity >= VERBOSITY_HIGH) || recordStats)
	{
		ostringstream strStream;
		strStream << tag.c_str() << "#: " << numSolutions << " (" << grammar.getInstanceSize() << ")";
		size_t objectiveCount = 0;
		size_t numStatSols = stopIndex-startIndex;
		for (size_t s = 0; s < parameters.objectives.stages(); s++)
		{
			for (size_t t = 0; t < parameters.objectives.stageSize(s); t++)
			{
				double measureMean = 0;
				for (size_t i = startIndex; i < stopIndex; i++)
				{
					measureMean += (*pSolutions)[i].result[s][t];
				}
				if (numStatSols)
				{
					measureMean /= numStatSols;
				}
				strStream << " " << parameters.objectives.objectiveName(s, t).c_str() << ": " << measureMean;
				if (recordStats)
				{
					double measureStd = 0;
					double measureMin = numeric_limits<double>::max();
					double measureMax = 0;
					for (size_t i = startIndex; i < stopIndex; i++)
					{
						double measureVal = (*pSolutions)[i].result[s][t];
						measureStd += (measureVal - measureMean) * (measureVal - measureMean);
						if (measureVal < measureMin) { measureMin = measureVal; }
						if (measureVal > measureMax) { measureMax = measureVal; }
					}
					if (numStatSols)
					{
						measureStd = sqrt(measureStd / numStatSols);
					}
					statMatrix[currentGeneration][objectiveCount++] = (*pSolutions)[startIndex].result[s][t];
					statMatrix[currentGeneration][objectiveCount++] = measureMean;
					statMatrix[currentGeneration][objectiveCount++] = measureStd;
					statMatrix[currentGeneration][objectiveCount++] = measureMin;
					statMatrix[currentGeneration][objectiveCount++] = measureMax;
				}
			}
		}
		globalMessageLog.printStream(COGMT_INFO_STATS, VERBOSITY_HIGH, strStream);
	}
}

void CogEvolver::computeGlobalStats(vector<vector<double>>& statMatrix)
{
	if (recordStats)
	{
		CogProductionStats prodStats;
		collectStats(prodStats, grammar.getPool());
		if (prodStats.instanceCount)
		{
			prodStats.instanceAgeMean /= prodStats.instanceCount;
			prodStats.instanceUsageMean /= prodStats.instanceCount;
		}
		collectStats2(prodStats, grammar.getPool());
		size_t prodSize = grammar.getInstanceSize();
		if (prodSize)
		{
			prodStats.instanceUsageStd = sqrt(prodStats.instanceUsageStd / prodSize);
			prodStats.instanceAgeStd = sqrt(prodStats.instanceAgeStd / prodSize);
		}
		statMatrix[currentGeneration][0] = currentGeneration+1;
		statMatrix[currentGeneration][3] = numSolutions;	
		statMatrix[currentGeneration][5] = prodSize;
		statMatrix[currentGeneration][6] = prodStats.instanceUsageMean;
		statMatrix[currentGeneration][7] = prodStats.instanceUsageStd;
		statMatrix[currentGeneration][8] = prodStats.instanceUsageMax;
		statMatrix[currentGeneration][9] = prodStats.instanceAgeMean;
		statMatrix[currentGeneration][10] = prodStats.instanceAgeStd;
		statMatrix[currentGeneration][11] = prodStats.instanceAgeMax;
		statMatrix[currentGeneration][12] = prodStats.ntSizeMean;
		statMatrix[currentGeneration][13] = prodStats.ntSizeStd;
		statMatrix[currentGeneration][14] = prodStats.ntSizeMin;
		statMatrix[currentGeneration][15] = prodStats.ntSizeMax;
	}
}

void CogEvolver::addClockStat(vector<vector<double>>& statMatrix, double time, size_t index)
{
	statMatrix[currentGeneration][index] += time;
}

void CogEvolver::addFailedStat(vector<vector<double>>& statMatrix, size_t invalidCount)
{
	statMatrix[currentGeneration][4] = invalidCount;
}

void CogEvolver::printSolutionStats(size_t index, const string& tag)
{
	if (parameters.verbosity >= VERBOSITY_NORMAL)
	{
		ostringstream strStream;
		strStream << tag.c_str() << resetiosflags(ios::fixed) << setprecision(6);
		for (size_t s = 0; s < parameters.objectives.stages(); s++)
		{
			for (size_t t = 0; t < parameters.objectives.stageSize(s); t++)
			{
				strStream << " " << parameters.objectives.objectiveName(s, t).c_str() << ": " << (*pSolutions)[index].result[s][t];
			}
		}
		globalMessageLog.printStream(COGMT_INFO_STATS, VERBOSITY_NORMAL, strStream);
	}
}

void CogEvolver::collectStats(CogProductionStats& stats, CogProductionPool* pPool)
{
	list<CogProduction>::iterator inst_iter = pPool->instances.begin();
	for (; inst_iter != pPool->instances.end(); ++inst_iter)
	{
		stats.instanceUsageMean += inst_iter->usedBy;
		stats.instanceAgeMean += currentGeneration - inst_iter->generation;
		stats.ntSizeMean += inst_iter->ntSize;
		stats.instanceCount++;
	}
	list<CogProductionPool>::iterator pool_iter = pPool->pools.begin();
	for (; pool_iter != pPool->pools.end(); ++pool_iter)
	{
		collectStats(stats, &*pool_iter);
	}
}

void CogEvolver::collectStats2(CogProductionStats& stats, CogProductionPool* pPool)
{
	list<CogProduction>::iterator inst_iter = pPool->instances.begin();
	for (; inst_iter != pPool->instances.end(); ++inst_iter)
	{
		stats.instanceUsageStd += (inst_iter->usedBy - stats.instanceUsageMean) * (inst_iter->usedBy - stats.instanceUsageMean);
		if (inst_iter->usedBy > stats.instanceUsageMax) { stats.instanceUsageMax = inst_iter->usedBy; }
		double instanceAge = currentGeneration - inst_iter->generation;
		stats.instanceAgeStd += (instanceAge  - stats.instanceAgeMean) * (instanceAge  - stats.instanceAgeMean);
		if (instanceAge > stats.instanceAgeMax) { stats.instanceAgeMax = instanceAge ; }
		stats.ntSizeStd += (inst_iter->ntSize - stats.ntSizeMean) * (inst_iter->ntSize - stats.ntSizeMean);
		if (inst_iter->ntSize < stats.ntSizeMin) { stats.ntSizeMin = inst_iter->ntSize; }
		if (inst_iter->ntSize > stats.ntSizeMax) { stats.ntSizeMax = inst_iter->ntSize; }
	}
	list<CogProductionPool>::iterator pool_iter = pPool->pools.begin();
	for (; pool_iter != pPool->pools.end(); ++pool_iter)
	{
		collectStats2(stats, &*pool_iter);
	}
}

// ------------------------------------------------------------------------------------------